import torch.nn as nn
import torch
from model.resnet import resnet18
from PIL import Image

def predict_raw(inputs):
    model_path = '/mnt/myproject/classfication/simple_classifier/saved/ckpt_epoch_49-21-08-29-15-58-34.pth'
    # device = torch.device('cuda:0')
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model_ft = resnet18(progress=True)
    num_ftrs = model_ft.fc.in_features
    # Here the size of each output sample is set to 2.
    # Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
    model_ft.fc = nn.Linear(num_ftrs, 2)

    checkpoint = torch.load(model_path, map_location='cpu')
    model_ft.load_state_dict(checkpoint['model'], strict=False)

    model_ft.to(device)
    with torch.no_grad():
        inputs.to(device)

        outputs = model_ft()
        _, preds = torch.max(outputs, 1)
    return preds



